/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: gradient Split Tune module model header file
 * Author: qianglihong
 * Create: 2020-09-15
 */

#ifndef AUTO_TUNING_SPLIT_MODEL_H
#define AUTO_TUNING_SPLIT_MODEL_H
#include <iostream>
#include <vector>
#include "layers.h"
#include "cluster.h"
#include "base.h"
#include "log.h"

constexpr uint32_t DEFAULT_BATCH_SIZE = 32;

class Model {
public:
    struct SliceMeth {
        std::vector<int> sliceRatio;
        std::vector<float> sliceSize;
    };

    Model(std::vector<uint64_t> &graInfoCost, std::vector<uint64_t> &graInfoSize, int batchs, u64 tensorLimit);

    ~Model();

    std::vector<int> GradientSlicing(Cluster &cluster, const Communication &op, int batchSize = DEFAULT_BATCH_SIZE);

private:
    float CalculateParamSize(std::vector<Layer> layers, int start = 0, int end = -1);

    float CalculateCost(std::vector<Layer> layers, int start = 0, int end = -1);

    int CalculateNextSlice(std::vector<Layer> &layers, int sliceStart, float costCommunication, float startUpCommTime);

    SliceMeth CalculateTrail(Cluster &cluster, const Communication &op, std::vector<Layer> &layersOriginal,
        int firstSliceEnd, float &trailCostNew);

    std::vector<Layer> mBps_;
    std::string mName_;
    int mBatchs_;
    float mLimit_;
};

#endif
